# @FileName  : test-hw.py
# @Time      : 2025/8/21 13:59
# @Author    : LuZhaoHui
# @Software  : PyCharm

import os
import time
import json
import requests
from PIL import Image
import io


class CustomVisionModel:
    def __init__(self, api_base_url, api_key, stream=False, model_id="vlm"):
        self.api_base_url = api_base_url.rstrip('/')
        self.api_key = api_key
        self.model_id = model_id
        self.stream = stream
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }

    def resize_image(self, image_path, max_length=768):
        """调整图片大小，确保长边不超过max_length，保持原比例"""
        with Image.open(image_path) as img:
            # 获取原图尺寸
            width, height = img.size

            # 计算缩放比例
            if width > height:
                # 宽为长边
                if width > max_length:
                    ratio = max_length / width
                    new_width = max_length
                    new_height = int(height * ratio)
                else:
                    new_width, new_height = width, height
            else:
                # 高为长边
                if height > max_length:
                    ratio = max_length / height
                    new_height = max_length
                    new_width = int(width * ratio)
                else:
                    new_width, new_height = width, height

            # 若无需缩放则直接返回原图字节流
            if (new_width, new_height) == (width, height):
                img_byte_arr = io.BytesIO()
                img.save(img_byte_arr, format='JPEG')
                return img_byte_arr.getvalue()

            # 按新尺寸缩放（使用高质量缩放算法）
            resized_img = img.resize((new_width, new_height), Image.LANCZOS)

            # 将缩放后的图片转为字节流
            img_byte_arr = io.BytesIO()
            resized_img.save(img_byte_arr, format='JPEG')
            return img_byte_arr.getvalue()

    def analyze_image_stream(self, image_path, prompt, max_tokens=10000, max_image_length=1920):
        try:
            if not os.path.exists(image_path):
                raise FileNotFoundError(f"图像文件不存在: {image_path}")

            # 处理图片（缩放并转为字节流）
            image_bytes = self.resize_image(image_path, max_image_length)

            # 编码处理后的图像
            import base64
            encoded_image = base64.b64encode(image_bytes).decode('utf-8')

            payload = {
                "model": self.model_id,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": f"data:image/jpeg;base64,{encoded_image}"
                                }
                            }
                        ]
                    }
                ],
                "max_tokens": max_tokens,
                "temperature": 0.0,
                "stream": self.stream  # 启用流式输出
            }

            # 发送流式请求
            with requests.post(
                    f"{self.api_base_url}/v1/chat/completions",
                    headers=self.headers,
                    data=json.dumps(payload),
                    stream=self.stream  # 保持连接打开，接收流式响应
            ) as response:
                if response.status_code != 200:
                    raise Exception(f"API调用失败: {response.status_code}, {response.text}")

                if self.stream:
                    # 逐行处理流式响应
                    for line in response.iter_lines():
                        if line:
                            # 记录当前块的接收时间
                            current_time = time.time()
                            # 去除数据前缀 "data: "
                            line = line.decode('utf-8').lstrip('data: ')
                            if line == '[DONE]':  # 流式结束标记
                                break
                            try:
                                chunk = json.loads(line)
                                # 提取内容
                                content = chunk.get("choices", [])[0].get("delta", {}).get("content", "")
                                if content:
                                    # 返回内容和当前时间戳
                                    yield (content, current_time)
                            except json.JSONDecodeError:
                                continue
                else:
                    return response.text
        except Exception as e:
            print(f"分析图像时出错: {str(e)}")
            return None


if __name__ == "__main__":
    API_BASE_URL = "http://192.168.132.92:3000"
    API_KEY = "xxx"
    isStream = True
    IMAGE_PATH = "F:/data/image/C/32-000741A.jpg"
    PROMPT = "这是一份人口登记卡片，第一步获取标题，第二步建立个人信息，获取上面的姓名、别名、性别、民族、出生日期、籍贯、出生地、职业、婚姻状况、现住址、登记日期。第三步建立户主信息，获取右上角户主姓名与户主关系信息。按照json中文格式数据。"
    # IMAGE_PATH = "F:/data/image/47.2-PH1.1-1-305/305-000126A.jpg"
    # PROMPT = "这是一份常住人口登记表，第一步获取标题内容和抬头的现住址内容。第二步根据表格列姓名判断总人数，按列姓名提取全部个人信息：姓名、别名、性别、出生日期、民族、籍贯、出生地、婚姻状况、文化程度、服务处所、职业、居民证编号。第三步判断户主，其他人与他的关系。按照JSON的中文格式输出"
    # IMAGE_PATH = "F:/data/image/1.1-73N+/bajiao/47.2-PH1.1-1-100/100-000002A.jpg"
    # PROMPT = "这是一份户口登记簿，是手写体，第一步获取标题内容和抬头的现住址内容。第二步根据表格列信息判断全部总人数，提取个人信息：姓名、性别、出生、民族、籍贯、出生地、文化程度、居民证编号。第三步判断户主，其他人与他的关系。按照JSON的中文格式输出"
    model = CustomVisionModel(
        api_base_url=API_BASE_URL,
        api_key=API_KEY,
        stream=isStream,
        model_id="vlm"
    )

    start_time = time.time()
    if isStream:
        print("开始分析，结果将流式输出...")
        first_chunk_time = None
        chunk_count = 0
        full_content = []

        # 处理流式输出
        for chunk, chunk_time in model.analyze_image_stream(IMAGE_PATH, PROMPT):
            # 记录第一个块的时间
            if first_chunk_time is None:
                first_chunk_time = chunk_time

            print(chunk, end='', flush=True)  # 实时打印，不换行
            full_content.append(chunk)
            chunk_count += 1

        print("\n\n分析完成")
        end_time = time.time()

        # 计算处理速度
        if chunk_count > 0 and first_chunk_time is not None:
            # 计算实际处理时间（从第一个块到最后一个块）
            processing_duration = end_time - first_chunk_time
            chunks_per_second = chunk_count / processing_duration if processing_duration > 0 else 0
            print(f"处理块数量: {chunk_count}个")
            print(f"处理速度: {chunks_per_second:.2f} tokens/秒")
    else:
        print("开始分析，结果输出...")

        response = model.analyze_image_stream(IMAGE_PATH, PROMPT)

        print(response)

        end_time = time.time()



    print(f"总耗时: {end_time - start_time:.2f}秒")
